import safetensors.torch

# 定义一个函数来替换权重名称
def replace_weight_names(input_file, output_file, name_mapping):
    # 从输入文件中加载 safetensor 数据
    tensors = safetensors.torch.load_file(input_file)
    new_tensors = {}
    # 遍历原始的权重名称和对应的张量
    for old_name, tensor in tensors.items():
        new_name = old_name
        for key,val in name_mapping.items():
            if key in old_name:
                new_name = old_name.replace(key, val)
                break
        # 将新名称和对应的张量添加到新的字典中
        new_tensors[new_name] = tensor
    # 将新的张量保存到输出文件中
    safetensors.torch.save_file(new_tensors, output_file)

# 示例使用
input_file = "model.safetensors"
output_file = "output.safetensors"
# 定义权重名称的映射表，键为旧名称，值为新名称
name_mapping = {
    "LayerNorm.gamma": "LayerNorm.weight",
    "LayerNorm.beta": "LayerNorm.bias"
}

# 调用函数进行权重名称替换
replace_weight_names(input_file, output_file, name_mapping)    
